#!/usr/bin/env python
# coding: utf-8

import logging
import os
import glob

import yaml
import matplotlib.pyplot as plt
import numpy as np
from datetime import timedelta
from tqdm import tqdm
from os.path import join
from time import sleep
import sys
sys.path = ['/home/hzh/.local/lib/python3.5/site-packages/'] + sys.path
from PIL import Image
# import cv2


# from models import load_pytorch_dense_net
from template_recognize_chi import TemplateMatcher

# predictor = load_pytorch_dense_net()
num_classes = 3753


def load_db(fp):
    with open(fp, 'r') as f:
        return yaml.load(f)

def main():
    db_index = load_db('./chi_seq_imgs/index.yaml')
    sbt_seq = []
    predictor = TemplateMatcher('./data/font_tmpls/')
    for k in tqdm(db_index.keys()):
        info = db_index[k]
        chi_chain = []
        for fn in info['chi_seq']:
            img = Image.open(join('./chi_seq_imgs', fn))
            img = np.asarray(img)
            c = predictor(img)[1]
            chi_chain.append(c)
        sbt_seq.append((info['time'], ''.join(chi_chain)))
    dump2srt(sbt_seq)

def time2str(t):
    dt = timedelta(seconds=t/1000.0)
    s = str(dt)
    return s.replace('.', ',')



def dump2srt(seq, fp="./a.srt"):
    lines = []
    for i in tqdm(range(len(seq))):
        lines.append("%d"%i)
        cur_tm = seq[i][0]
        next_tm = seq[i+1][0] if i+1 < len(seq) else cur_tm+1000 # plus 1 seconds
        lines.append("%s --> %s" % (time2str(cur_tm), time2str(next_tm)) )
        lines.append("%s\n" % seq[i][1])
    with open(fp, 'w') as f:
        f.write('\n'.join(lines))



if '__main__' == __name__:
    main()
